# -*- coding: utf-8 -*-
"""
ATR (Average True Range) 分析工具
用于计算股票的ATR指标、波动率，并分析支撑位和压力位
"""
import akshare as ak
import pandas as pd
import numpy as np
import datetime
import os
import argparse
import time
import random
from functools import lru_cache

# 导入安全请求模块
from akshare_request import (
    get_stock_zh_a_hist_safe,
    get_stock_zh_a_spot_em_safe,
    get_tool_trade_date_hist_sina_safe
)

# 配置参数
ANALYSIS_CONFIG = {
    'short_term_days': 5,      # 短线天数
    'trading_days': 30,        # 查询最近的交易日数量
    'atr_period': 14,          # ATR计算周期
    'min_market_value': 1000000000,  # 最小市值（10亿）
    'max_market_value': 500000000000,  # 最大市值（5000亿）
    'include_st': False,       # 是否包含ST股票
    'include_gem': False,      # 是否包含创业板（30开头）
    'include_sci_tech': False, # 是否包含科创板（68开头）
    'include_main': True,      # 是否包含主板（00、60开头）
    'output_folder': 'data',   # 输出文件夹
    # 'stock_codes': ['601919','600989','600219','000792','002709', '601888','600150','002600','001206','603606']  # 默认分析的股票代码
    'stock_codes': ['600177','002170','002895','600096','600989','000792','002570','601919']  # 默认分析的股票代码
}

def format_date(date_str):
    """格式化日期字符串为YYYYMMDD格式"""
    # 如果输入的是YYYY-MM-DD格式，转换为YYYYMMDD
    if '-' in date_str:
        return date_str.replace('-', '')
    return date_str

def get_trading_dates(days=30):
    """获取最近N个交易日的日期列表，使用安全请求函数"""
    try:
        # 获取当前日期
        today = datetime.datetime.now().date()  # 转换为date对象
        # 计算60天前的日期（为了确保能获取到足够的交易日）
        days_ago_60 = today - datetime.timedelta(days=60)
        
        # 使用安全请求获取交易日历
        trade_cal = get_tool_trade_date_hist_sina_safe()
        # 筛选指定日期范围内的交易日
        trade_dates = trade_cal[(trade_cal['trade_date'] >= days_ago_60) & 
                               (trade_cal['trade_date'] <= today)]
        
        # 获取最近N个交易日
        recent_dates = trade_dates.tail(days)['trade_date'].tolist()
        
        # 确保日期格式一致 (YYYYMMDD) - 转换为akshare的stock_zh_a_hist函数所需格式
        formatted_dates = []
        for date in recent_dates:
            # 将datetime.date对象转换为YYYYMMDD格式字符串
            date_str = date.strftime('%Y%m%d')
            formatted_dates.append(date_str)
        
        return formatted_dates
    except Exception as e:
        print(f"获取交易日期失败: {e}")
        return []

@lru_cache(maxsize=300)
def get_stock_history(stock_code, start_date, end_date, max_retries=3):
    """获取指定股票在给定时间范围内的历史数据，使用安全请求函数"""
    try:
        # 格式化日期
        start_date = format_date(start_date)
        end_date = format_date(end_date)
        
        # 使用安全请求函数获取历史数据
        # 该函数内部已包含重试、限流和缓存机制
        df = get_stock_zh_a_hist_safe(
            symbol=stock_code, 
            period="daily", 
            start_date=start_date, 
            end_date=end_date, 
            adjust="qfq"
        )
        
        return df if not df.empty else pd.DataFrame()
    except Exception as e:
        print(f"{stock_code}: 获取历史数据失败 {e}")
        return pd.DataFrame()

def calculate_atr(df, period=14):
    """计算ATR (Average True Range)"""
    try:
        # 确保数据按日期排序
        if '日期' in df.columns:
            df = df.sort_values('日期')
        elif 'trade_date' in df.columns:
            df = df.sort_values('trade_date')
        
        # 计算真实波幅 (True Range)
        # TR = max(high - low, abs(high - prev_close), abs(low - prev_close))
        df['prev_close'] = df['收盘'].shift(1)
        df['TR1'] = df['最高'] - df['最低']
        df['TR2'] = abs(df['最高'] - df['prev_close'])
        df['TR3'] = abs(df['最低'] - df['prev_close'])
        df['TR'] = df[['TR1', 'TR2', 'TR3']].max(axis=1)
        
        # 计算ATR (使用简单移动平均)
        df['ATR'] = df['TR'].rolling(window=period).mean()
        
        # 清理临时列
        df = df.drop(['TR1', 'TR2', 'TR3', 'prev_close'], axis=1)
        
        return df
    except Exception as e:
        print(f"计算ATR失败: {e}")
        return df

def calculate_volatility(df, short_term_days=3):
    """计算短期波动率"""
    try:
        # 确保数据按日期排序
        if '日期' in df.columns:
            df = df.sort_values('日期')
        elif 'trade_date' in df.columns:
            df = df.sort_values('trade_date')
        
        # 计算日涨跌幅
        if '涨跌幅' in df.columns:
            df['daily_change'] = df['涨跌幅']
        elif '涨跌幅(%)' in df.columns:
            df['daily_change'] = df['涨跌幅(%)']
        else:
            df['daily_change'] = df['收盘'].pct_change() * 100
        
        # 计算短期平均波动率
        df['volatility'] = df['daily_change'].rolling(window=short_term_days).mean()
        
        # 计算短期波动率标准差
        df['volatility_std'] = df['daily_change'].rolling(window=short_term_days).std()
        
        return df
    except Exception as e:
        print(f"计算波动率失败: {e}")
        return df

def calculate_support_resistance(df, atr_multiplier=2):
    """计算支撑位和压力位"""
    try:
        # 使用最新收盘价和ATR计算支撑位和压力位
        latest_close = df['收盘'].iloc[-1]
        latest_atr = df['ATR'].iloc[-1]
        
        # 支撑位 = 收盘价 - ATR倍数 * ATR
        support = latest_close - atr_multiplier * latest_atr
        
        # 压力位 = 收盘价 + ATR倍数 * ATR
        resistance = latest_close + atr_multiplier * latest_atr
        
        return support, resistance
    except Exception as e:
        print(f"计算支撑位和压力位失败: {e}")
        return None, None

def analyze_stock_volatility(stock_code, stock_name=None, trading_dates=None):
    """分析指定股票的波动率和ATR"""
    try:
        if not trading_dates or len(trading_dates) < ANALYSIS_CONFIG['atr_period']:
            print(f"交易日数据不足，无法计算ATR")
            return None
            
        # 获取开始日期和结束日期
        start_date = trading_dates[0]
        end_date = trading_dates[-1]
        
        # 获取历史数据
        df = get_stock_history(stock_code, start_date, end_date)
        if df.empty:
            print(f"{stock_code}: 获取历史数据失败")
            return None
        
        # 计算ATR
        df = calculate_atr(df, period=ANALYSIS_CONFIG['atr_period'])
        
        # 计算波动率
        df = calculate_volatility(df, short_term_days=ANALYSIS_CONFIG['short_term_days'])
        
        # 计算支撑位和压力位
        support, resistance = calculate_support_resistance(df)
        
        # 获取最新数据
        latest_data = df.iloc[-1]
        latest_atr = latest_data['ATR']
        latest_volatility = latest_data['volatility']
        latest_volatility_std = latest_data['volatility_std']
        latest_close = latest_data['收盘']
        
        # 计算正常波动范围
        normal_low = latest_volatility - latest_volatility_std
        normal_high = latest_volatility + latest_volatility_std
        
        # 判断极端波动阈值 (平均波动率 + 2倍标准差)
        extreme_threshold = latest_volatility + 2 * latest_volatility_std
        
        # 生成分析结论
        if latest_volatility >= extreme_threshold:
            conclusion = "异常波动，建议谨慎操作"
        elif latest_volatility > normal_high:
            conclusion = "波动偏大，注意风险控制"
        elif latest_volatility < normal_low:
            conclusion = "波动偏小，可能处于盘整阶段"
        else:
            conclusion = "波动正常，可按既定策略操作"
        
        # 获取股票名称（如果未提供）
        if not stock_name:
            try:
                stock_info = get_stock_zh_a_spot_em_safe()
                stock_info = stock_info[stock_info['代码'] == stock_code]
                if not stock_info.empty:
                    stock_name = stock_info.iloc[0]['名称']
                else:
                    stock_name = "未知"
            except:
                stock_name = "未知"
        
        # 返回分析结果
        return {
            "代码": stock_code,
            "名称": stock_name,
            "收盘价": round(latest_close, 2),
            "ATR(周期={})".format(ANALYSIS_CONFIG['atr_period']): round(latest_atr, 2),
            "{}日平均波动率(%)".format(ANALYSIS_CONFIG['short_term_days']): round(latest_volatility, 2),
            "波动率标准差(%)": round(latest_volatility_std, 2),
            "正常波动范围(%)": f"{round(normal_low, 2)} ~ {round(normal_high, 2)}",
            "极端波动阈值(%)": round(extreme_threshold, 2),
            "支撑位": round(support, 2),
            "压力位": round(resistance, 2),
            "分析结论": conclusion
        }
    except Exception as e:
        print(f"{stock_code}: 分析波动率失败 {e}")
        return None

def analyze_stocks(stock_codes=None):
    """分析指定股票列表的波动率和ATR"""
    # 获取交易日期
    trading_dates = get_trading_dates(ANALYSIS_CONFIG['trading_days'])
    if not trading_dates:
        print("获取交易日期失败，程序结束")
        return pd.DataFrame()
        
    # 格式化显示日期（YYYY-MM-DD）
    start_display = trading_dates[0][:4] + '-' + trading_dates[0][4:6] + '-' + trading_dates[0][6:] if len(trading_dates[0]) == 8 else trading_dates[0]
    end_display = trading_dates[-1][:4] + '-' + trading_dates[-1][4:6] + '-' + trading_dates[-1][6:] if len(trading_dates[-1]) == 8 else trading_dates[-1]
    print(f"\n===== 分析最近{ANALYSIS_CONFIG['trading_days']}个交易日（{start_display}至{end_display}）波动率情况 =====\n")
    
    # 如果未提供股票代码，使用配置中的默认代码
    if not stock_codes:
        stock_codes = ANALYSIS_CONFIG['stock_codes']
    
    # 分析每只股票
    results = []
    for stock_code in stock_codes:
        print(f"分析股票: {stock_code}")
        result = analyze_stock_volatility(stock_code, trading_dates=trading_dates)
        if result:
            results.append(result)
            
            # 打印单只股票的分析结果
            print(f"\n股票: {result['名称']}({result['代码']})")
            print(f"收盘价: {result['收盘价']}")
            print(f"ATR(周期={ANALYSIS_CONFIG['atr_period']}): {result['ATR(周期={})'.format(ANALYSIS_CONFIG['atr_period'])]}")
            print(f"{ANALYSIS_CONFIG['short_term_days']}日平均波动率: {result['{}日平均波动率(%)'.format(ANALYSIS_CONFIG['short_term_days'])]}%")
            print(f"波动率标准差: {result['波动率标准差(%)']}%")
            print(f"正常波动范围: {result['正常波动范围(%)']}")
            print(f"极端波动阈值: {result['极端波动阈值(%)']}%")
            print(f"支撑位: {result['支撑位']}")
            print(f"压力位: {result['压力位']}")
            print(f"分析结论: {result['分析结论']}")
            print("-" * 50)
    
    # 转换为DataFrame
    if not results:
        print("没有有效的分析结果")
        return pd.DataFrame()
    
    df_results = pd.DataFrame(results)
    return df_results

def save_results(df):
    """保存结果到文件（Markdown格式）"""
    if df.empty:
        return
    
    # 创建输出文件夹
    folder_path = os.path.join(os.path.dirname(__file__), ANALYSIS_CONFIG['output_folder'])
    os.makedirs(folder_path, exist_ok=True)
    
    # 获取当前日期
    today = datetime.datetime.now().strftime('%Y%m%d')
    
    # 生成文件名
    file_name = f"波动率分析_{ANALYSIS_CONFIG['short_term_days']}日_{today}.md"
    file_path = os.path.join(folder_path, file_name)
    
    # 生成Markdown内容
    md_content = f"# 股票波动率分析报告\n\n"
    md_content += f"## 分析参数\n\n"
    md_content += f"- 分析日期: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
    md_content += f"- 短线天数: {ANALYSIS_CONFIG['short_term_days']}\n"
    md_content += f"- 交易日数量: {ANALYSIS_CONFIG['trading_days']}\n"
    md_content += f"- ATR计算周期: {ANALYSIS_CONFIG['atr_period']}\n\n"
    
    md_content += f"## 分析结果\n\n"
    
    # 为每只股票生成详细分析
    for _, row in df.iterrows():
        md_content += f"### {row['名称']}({row['代码']})\n\n"
        md_content += f"- 收盘价: {row['收盘价']}\n"
        md_content += f"- ATR(周期={ANALYSIS_CONFIG['atr_period']}): {row['ATR(周期={})'.format(ANALYSIS_CONFIG['atr_period'])]}\n"
        md_content += f"- {ANALYSIS_CONFIG['short_term_days']}日平均波动率: {row['{}日平均波动率(%)'.format(ANALYSIS_CONFIG['short_term_days'])]}%\n"
        md_content += f"- 波动率标准差: {row['波动率标准差(%)']}%\n"
        md_content += f"- 正常波动范围: {row['正常波动范围(%)']}\n"
        md_content += f"- 极端波动阈值: {row['极端波动阈值(%)']}%\n"
        md_content += f"- 支撑位: {row['支撑位']}\n"
        md_content += f"- 压力位: {row['压力位']}\n"
        md_content += f"- **分析结论**: {row['分析结论']}\n\n"
        
        # 添加波动率评估图表（文本形式）
        volatility = row['{}日平均波动率(%)'.format(ANALYSIS_CONFIG['short_term_days'])]
        std = row['波动率标准差(%)']
        extreme = row['极端波动阈值(%)']
        
        md_content += "#### 波动率评估\n\n"
        md_content += "```\n"
        md_content += "低波动 <----|----|----|----|----> 高波动\n"
        
        # 计算波动率在图表中的位置
        position = int((volatility + 5) * 2)  # 假设范围从-5%到+5%
        position = max(0, min(position, 20))  # 限制在0-20的范围内
        
        chart = list("                     ")
        chart[position] = "▲"
        
        # 标记正常范围和极端阈值
        normal_low_pos = int((volatility - std + 5) * 2)
        normal_low_pos = max(0, min(normal_low_pos, 20))
        
        normal_high_pos = int((volatility + std + 5) * 2)
        normal_high_pos = max(0, min(normal_high_pos, 20))
        
        extreme_pos = int((extreme + 5) * 2)
        extreme_pos = max(0, min(extreme_pos, 20))
        
        chart[normal_low_pos] = "["
        chart[normal_high_pos] = "]"
        if extreme_pos < 20:
            chart[extreme_pos] = "!"
        
        md_content += "".join(chart) + "\n"
        md_content += "```\n\n"
        md_content += f"- [ ] 表示正常波动范围\n"
        md_content += f"- ! 表示极端波动阈值\n"
        md_content += f"- ▲ 表示当前波动率\n\n"
        
        md_content += "---\n\n"
    
    # 添加总结
    md_content += "## 总结\n\n"
    md_content += f"本次分析共包含 {len(df)} 只股票。\n\n"
    
    # 统计不同结论的股票数量
    conclusion_counts = df['分析结论'].value_counts()
    md_content += "### 结论分布\n\n"
    for conclusion, count in conclusion_counts.items():
        md_content += f"- {conclusion}: {count}只\n"
    
    # 写入文件
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(md_content)
    
    print(f"\n数据已保存至 {file_path}")

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="分析股票的ATR和波动率")
    parser.add_argument("-s", "--short-term", type=int, help="短线天数")
    parser.add_argument("-d", "--days", type=int, help="查询最近的交易日数量")
    parser.add_argument("-a", "--atr-period", type=int, help="ATR计算周期")
    parser.add_argument("--min-value", type=float, help="最小市值（亿元）")
    parser.add_argument("--max-value", type=float, help="最大市值（亿元）")
    parser.add_argument("--include-st", action="store_true", help="是否包含ST股票")
    parser.add_argument("--include-gem", action="store_true", help="是否包含创业板")
    parser.add_argument("--include-sci-tech", action="store_true", help="是否包含科创板")
    parser.add_argument("--exclude-main", action="store_true", help="是否排除主板")
    parser.add_argument("-o", "--output", help="输出文件夹")
    parser.add_argument("-c", "--codes", nargs='+', help="要分析的股票代码列表")
    
    args = parser.parse_args()
    
    # 更新配置
    if args.short_term:
        ANALYSIS_CONFIG['short_term_days'] = args.short_term
    if args.days:
        ANALYSIS_CONFIG['trading_days'] = args.days
    if args.atr_period:
        ANALYSIS_CONFIG['atr_period'] = args.atr_period
    if args.min_value:
        ANALYSIS_CONFIG['min_market_value'] = args.min_value * 100000000  # 亿元转换为元
    if args.max_value:
        ANALYSIS_CONFIG['max_market_value'] = args.max_value * 100000000  # 亿元转换为元
    if args.include_st:
        ANALYSIS_CONFIG['include_st'] = True
    if args.include_gem:
        ANALYSIS_CONFIG['include_gem'] = True
    if args.include_sci_tech:
        ANALYSIS_CONFIG['include_sci_tech'] = True
    if args.exclude_main:
        ANALYSIS_CONFIG['include_main'] = False
    if args.output:
        ANALYSIS_CONFIG['output_folder'] = args.output
    if args.codes:
        ANALYSIS_CONFIG['stock_codes'] = args.codes

def main():
    """主函数"""
    # 解析命令行参数
    parse_arguments()
    
    start_time = time.time()
    print(f"开始分析: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"分析最近 {ANALYSIS_CONFIG['trading_days']} 个交易日的波动率情况")
    print(f"短线天数: {ANALYSIS_CONFIG['short_term_days']}")
    print(f"ATR计算周期: {ANALYSIS_CONFIG['atr_period']}")
    print(f"分析股票: {', '.join(ANALYSIS_CONFIG['stock_codes'])}")
    
    try:
        # 分析股票
        results = analyze_stocks()
        
        # 保存结果
        if not results.empty:
            save_results(results)
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
    finally:
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"\n分析完成，耗时: {elapsed_time:.2f} 秒 ({elapsed_time/60:.2f} 分钟)")
        print(f"结束时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")


if __name__ == "__main__":
    main()